{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 自定义 LocalTool\n",
    "\n",
    "LocalTool 是能够将代码运行在开发者本地的模块，允许开发者自定义模块功能，不需要任何第三方的远程服务即可实现自定义工具、对接公司内部业务逻辑等模块。\n",
    "\n",
    "在本章节中将通过创建一个单词本为例，介绍如何使用 LocalTool 模块。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. AddWordTool\n",
    "\n",
    "### 1.1 需求描述\n",
    "\n",
    "假设我们要创建一个单词本的功能 ，其中最主要的功能是：添加单词。\n",
    "\n",
    "* 输入参数：单词\n",
    "* 输出参数：是否添加成功，并返回对应的描述信息\n",
    "* 单词本的功能描述：描述单词本的主要功能，也是智能体（Agent）从用户的 query 中判别对应意图的关键描述。\n",
    "\n",
    "### 1.2 代码实现\n",
    "\n",
    "创建一个 LocalTool 需要几个关键信息，详细可参考：[Tool 关键信息](/modules/tools/#31)，对应代码模块需要实现：\n",
    "\n",
    "* 创建 Tool 输入和输出参数的类\n",
    "* 继承 Tool 并构建派生类\n",
    "* 实现 Tool 中的 `__call__`方法，实现主要的添加逻辑。\n",
    "\n",
    "\n",
    "具体代码实现如下所示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "## 添加环境环境变量\n",
    "import os\n",
    "\n",
    "os.environ[\"EB_AGENT_ACCESS_TOKEN\"] = \"<access_token>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AgentResponse(text='要将单词“red”添加到词库中，您需要执行以下步骤：\\n\\n1. 打开词库文件：找到您要添加单词的词库文件，并使用文本编辑器打开它。\\n2. 添加单词：在词库文件中找到要插入“red”单词的位置，并在该位置插入该单词。确保单词的拼写和格式正确。\\n3. 保存文件：保存对词库文件的更改。\\n\\n请注意，具体的步骤可能会因您使用的词库软件而有所不同。如果您使用的是特定的词库软件或框架，请参考该软件或框架的文档以获取更详细的指导。', chat_history=[<HumanMessage role: 'user', content: '将单词：“red”添加到词库当中', token_count: 113>, <AIMessage role: 'assistant', function_call: {'name': 'AddWordTool', 'thoughts': '用户希望将单词添加到词库当中，我可以使用AddWordTool工具来完成这个任务。', 'arguments': '{\"word\":\"red\"}'}, token_count: 28>, <FunctionMessage role: 'function', name: 'AddWordTool', content: '{\"result\": \"<red>单词已添加成功, 当前单词本中有如下单词：red\"}'>, <AIMessage role: 'assistant', content: '要将单词“red”添加到词库中，您需要执行以下步骤：\\n\\n1. 打开词库文件：找到您要添加单词的词库文件，并使用文本编辑器打开它。\\n2. 添加单词：在词库文件中找到要插入“red”单词的位置，并在该位置插入该单词。确保单词的拼写和格式正确。\\n3. 保存文件：保存对词库文件的更改。\\n\\n请注意，具体的步骤可能会因您使用的词库软件而有所不同。如果您使用的是特定的词库软件或框架，请参考该软件或框架的文档以获取更详细的指导。', token_count: 136>], steps=[ToolStep(info={'tool_name': 'AddWordTool', 'tool_args': '{\"word\":\"red\"}'}, result='{\"result\": \"<red>单词已添加成功, 当前单词本中有如下单词：red\"}', input_files=[], output_files=[])], status='FINISHED')\n"
     ]
    }
   ],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from typing import Any, Dict, Type, List\n",
    "from pydantic import Field\n",
    "from erniebot_agent.tools.base import Tool\n",
    "\n",
    "from erniebot_agent.tools.schema import ToolParameterView\n",
    "\n",
    "from erniebot_agent.agents.function_agent import FunctionAgent\n",
    "from erniebot_agent.chat_models import ERNIEBot\n",
    "from erniebot_agent.memory import WholeMemory\n",
    "\n",
    "\n",
    "class AddWordInput(ToolParameterView):\n",
    "    word: str = Field(description=\"待添加的单词\")\n",
    "\n",
    "\n",
    "class AddWordOutput(ToolParameterView):\n",
    "    result: str = Field(description=\"表示是否成功将单词成功添加到词库当中\")\n",
    "\n",
    "\n",
    "class AddWordTool(Tool):\n",
    "    description: str = \"添加单词到词库当中\"\n",
    "    input_type: Type[ToolParameterView] = AddWordInput\n",
    "    ouptut_type: Type[ToolParameterView] = AddWordOutput\n",
    "\n",
    "    def __init__(self) -> None:\n",
    "        self.word_books = {}\n",
    "        super().__init__()\n",
    "\n",
    "    async def __call__(self, word: str) -> Dict[str, Any]:\n",
    "        if word in self.word_books:\n",
    "            return {\"result\": f\"<{word}>单词已经存在，无需添加\"}\n",
    "        self.word_books[word] = True\n",
    "        words = \"\\n\".join(list(self.word_books.keys()))\n",
    "        return {\"result\": f\"<{word}>单词已添加成功, 当前单词本中有如下单词：{words}\"}\n",
    "\n",
    "\n",
    "agent = FunctionAgent(ERNIEBot(\"ernie-3.5\"), tools=[AddWordTool()], memory=WholeMemory())\n",
    "result = await agent.run(\"将单词：“red”添加到词库当中\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 代码详解\n",
    "\n",
    "以上代码整体逻辑为：\n",
    "\n",
    "1. 创建 Tool 所需的输入和输出参数的类型注释类，用于校验智能体在从 query 中抽取的参数。\n",
    "2. 创建 Tool 派生类，首先整体逻辑，此时需要注意 `__call__` 的输入参数需要与 `input_type` 中定义的参数列表一致，这样才能够正确接受参数。\n",
    "3. 将 `llm` 和 `tools` 都作为属性传递到 `FunctionAgent` 当中，此智能体能够处理 tools 的智能编排。\n",
    "\n",
    "AgentResponse 包含以下几个属性：\n",
    "* text: 智能体当前对话的最终输出\n",
    "* chat_history: 智能体在与 用户和 Tool 交互的过程中产生的所有消息列表。\n",
    "* files: 智能体与用户交互过程中产生的文件列表。\n",
    "* input_files: 智能体与用户交互过程中用户输入的文件列表。\n",
    "* output_files: 智能体与用户交互过程中工具输出的文件列表。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此案例中介绍了如何将单词添加到一个程序变量中，可是这个并不适用于实际场景，接下来将介绍如何将 Tool 和 外部文件进行联动。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. AddWordToJsonFileTool\n",
    "\n",
    "此案例将会介绍如何将单词添加到本地 Json 文件当中，从而实现持久化保存单词的功能\n",
    "\n",
    "### 2.1 需求分析\n",
    "\n",
    "此功能相比以上代码需要改变的仅仅只是将保存的目标改为Json 文件，调整包含：\n",
    "\n",
    "1. 工具初始化时，需要传入目标存储文件路径。\n",
    "2. 工具 `__call__` 方法中需要修改保存逻辑。\n",
    "\n",
    "代码实现如下所示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "\n",
    "def write_word_book(word_book: dict, file: str):\n",
    "    with open(file, \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump(word_book, f, ensure_ascii=False, indent=8)\n",
    "\n",
    "\n",
    "def read_word_book(file: str):\n",
    "    if not os.path.exists(file):\n",
    "        write_word_book({}, file)\n",
    "    with open(file, \"r\", encoding=\"utf-8\") as f:\n",
    "        data = json.load(f)\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AddWordTool(Tool):\n",
    "    description: str = \"添加单词到词库当中\"\n",
    "    input_type: Type[ToolParameterView] = AddWordInput\n",
    "    ouptut_type: Type[ToolParameterView] = AddWordOutput\n",
    "\n",
    "    def __init__(self, file: str):\n",
    "        self.file = file\n",
    "\n",
    "    async def __call__(self, word: str) -> Dict[str, Any]:\n",
    "        word_book = read_word_book(self.file)\n",
    "        if word in word_book:\n",
    "            return {\"result\": f\"<{word}>单词已经存在，无需添加\"}\n",
    "\n",
    "        word_book[word] = \"true\"\n",
    "        write_word_book(word_book, self.file)\n",
    "        return {\"result\": f\"<{word}>单词已添加成功\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AgentResponse(text='在Python中，你可以使用字典来创建一个词库，其中键是单词，值是单词的含义。以下是如何将单词“red”添加到词库的示例：\\n\\n\\n```python\\n# 创建一个空的字典\\nword_bank = {}\\n\\n# 将单词“red”添加到词库中\\nword_bank[\"red\"] = \"颜色\"\\n\\n# 打印词库\\nprint(word_bank)\\n```\\n这将在控制台上输出：\\n\\n\\n```arduino\\n{\\'red\\': \\'颜色\\'}\\n```\\n现在，你可以通过键“red”从词库中检索单词的含义，即“颜色”。', chat_history=[<HumanMessage role: 'user', content: '将单词：“red”添加到词库当中', token_count: 105>, <AIMessage role: 'assistant', function_call: {'name': 'AddWordTool', 'thoughts': '用户希望将单词添加到词库当中，我可以使用AddWordTool工具来完成这个任务。', 'arguments': '{\"word\":\"red\"}'}, token_count: 28>, <FunctionMessage role: 'function', name: 'AddWordTool', content: '{\"result\": \"<red>单词已经存在，无需添加\"}'>, <AIMessage role: 'assistant', content: '在Python中，你可以使用字典来创建一个词库，其中键是单词，值是单词的含义。以下是如何将单词“red”添加到词库的示例：\\n\\n\\n```python\\n# 创建一个空的字典\\nword_bank = {}\\n\\n# 将单词“red”添加到词库中\\nword_bank[\"red\"] = \"颜色\"\\n\\n# 打印词库\\nprint(word_bank)\\n```\\n这将在控制台上输出：\\n\\n\\n```arduino\\n{\\'red\\': \\'颜色\\'}\\n```\\n现在，你可以通过键“red”从词库中检索单词的含义，即“颜色”。', token_count: 136>], steps=[ToolStep(info={'tool_name': 'AddWordTool', 'tool_args': '{\"word\":\"red\"}'}, result='{\"result\": \"<red>单词已经存在，无需添加\"}', input_files=[], output_files=[])], status='FINISHED')\n",
      "AgentResponse(text='要在已有的词库中添加单词“blue”，你可以直接使用相同的键值对添加到字典中。以下是如何做到这一点的示例：\\n\\n\\n```python\\n# 已存在的词库\\nword_bank = {\\n    \"red\": \"颜色\"\\n}\\n\\n# 将单词“blue”添加到词库中\\nword_bank[\"blue\"] = \"颜色\"\\n\\n# 打印更新后的词库\\nprint(word_bank)\\n```\\n这将在控制台上输出：\\n\\n\\n```arduino\\n{\\'red\\': \\'颜色\\', \\'blue\\': \\'颜色\\'}\\n```\\n现在，词库包含两个单词：“red”和“blue”，它们都表示“颜色”。', chat_history=[<HumanMessage role: 'user', content: '将单词：“blue”添加到词库当中', token_count: 251>, <AIMessage role: 'assistant', function_call: {'name': 'AddWordTool', 'thoughts': '用户想要将单词添加到词库当中，我可以调用AddWordTool工具来完成这个任务。', 'arguments': '{\"word\":\"blue\"}'}, token_count: 28>, <FunctionMessage role: 'function', name: 'AddWordTool', content: '{\"result\": \"<blue>单词已经存在，无需添加\"}'>, <AIMessage role: 'assistant', content: '要在已有的词库中添加单词“blue”，你可以直接使用相同的键值对添加到字典中。以下是如何做到这一点的示例：\\n\\n\\n```python\\n# 已存在的词库\\nword_bank = {\\n    \"red\": \"颜色\"\\n}\\n\\n# 将单词“blue”添加到词库中\\nword_bank[\"blue\"] = \"颜色\"\\n\\n# 打印更新后的词库\\nprint(word_bank)\\n```\\n这将在控制台上输出：\\n\\n\\n```arduino\\n{\\'red\\': \\'颜色\\', \\'blue\\': \\'颜色\\'}\\n```\\n现在，词库包含两个单词：“red”和“blue”，它们都表示“颜色”。', token_count: 148>], steps=[ToolStep(info={'tool_name': 'AddWordTool', 'tool_args': '{\"word\":\"blue\"}'}, result='{\"result\": \"<blue>单词已经存在，无需添加\"}', input_files=[], output_files=[])], status='FINISHED')\n",
      "本地 words.json 文件的内容如下：\n",
      "{\n",
      "        \"red\": \"true\",\n",
      "        \"blue\": \"true\"\n",
      "}"
     ]
    }
   ],
   "source": [
    "agent = FunctionAgent(ERNIEBot(\"ernie-3.5\"), tools=[AddWordTool(\"words.json\")], memory=WholeMemory())\n",
    "result = await agent.run(\"将单词：“red”添加到词库当中\")\n",
    "print(result)\n",
    "result = await agent.run(\"将单词：“blue”添加到词库当中\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo \"本地 words.json 文件的内容如下：\"\n",
    "!cat words.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在实际场景中，通常不会将数据保存到本地文件当中，而是数据库当中，所以此时开发者只需要将对应的保存逻辑调整成数据库的相关逻辑代码即可。\n",
    "\n",
    "此案例是为了展示Tool 的扩展性，开发者可在输入输出参数、Tool 的代码实现上灵活扩展，实现自己的业务逻辑，比如：调用外部 HTTP API、内置 LLM 实现自定义 ReAct、Tool 嵌套 Tool（站在巨人的肩膀上）甚至 Tool 调用外部 Agent 实现更上层的功能。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
